# coding: utf-8
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from block import DBL



class yolov3_tiny_backbone(nn.Module):
    def __init__(self, in_ch=3, base_ch=16):
        super(yolov3_tiny_backbone, self).__init__()

        # 1
        self.conv1 = DBL(in_ch, base_ch)
        in_ch = base_ch
        base_ch *= 2

        # 2
        self.conv2 = DBL(in_ch, base_ch)
        in_ch = base_ch
        base_ch *= 2

        # 3
        self.conv3 = DBL(in_ch, base_ch)
        in_ch = base_ch
        base_ch *= 2

        # 4
        self.conv4 = DBL(in_ch, base_ch)
        in_ch = base_ch
        base_ch *= 2

        # 5
        self.conv5 = DBL(in_ch, base_ch)
        in_ch = base_ch
        base_ch *= 2

        # 6
        self.conv6 = DBL(in_ch, base_ch)
        in_ch = base_ch
        base_ch *= 2

        self.conv7 = DBL(in_ch, base_ch)


    def forward(self, x):
            x = F.max_pool2d(self.conv1(x), 2, 2)
            x = F.max_pool2d(self.conv2(x), 2, 2)
            x = F.max_pool2d(self.conv3(x), 2, 2)
            x = F.max_pool2d(self.conv4(x), 2, 2)
            x = F.max_pool2d(self.conv5(x), 2, 2)
            x1 = self.conv6(x)
            x = F.max_pool2d(x1, 2, 1, 1)
            x = self.conv7(x)
            return x1, x


class yolov3_tiny(nn.Module):
    def __init__(self, in_ch=3, base_ch=16):
        super(yolov3_tiny, self).__init__()
        self.backbone = yolov3_tiny_backbone(in_ch, base_ch)

        # head 1
        ch = base_ch * 64
        self.head1_1 = DBL(ch, 256, kernel=1, stride=1, pad=0)
        self.head1_2 = DBL(256, 512)
        self.head1_3 = nn.Conv2d(512, 255, kernel_size=1, stride=1, padding=0)

        # head 2
        self.head2_1 = DBL(256, 128, kernel=1, stride=1, pad=0)
        ch = base_ch * 32
        self.head2_2 = DBL(128+ch, 256)
        self.head2_3 = nn.Conv2d(256, 255, kernel_size=1, stride=1, padding=0)



    def forward(self, x):
        x2, x1 = self.backbone(x)

        # head 1
        x11 = self.head1_1(x1)
        x1 = self.head1_2(x11)
        y1 = self.head1_3(x1)

        # head 2
        x21 = self.head2_1(x11)
        x21 = F.interpolate(x21, scale_factor=2.0)
        x2 = torch.cat([x2, x21], dim=1)
        x2 = self.head2_2(x2)
        y2 = self.head2_3(x2)

        return y1, y2


if __name__ == "__main__":
    x = torch.randn(1,3,416, 416)
    m = yolov3_tiny(3, 16)
    y = m(x)
    print(y.shape)
